from langchain_core.example_selectors import LengthBasedExampleSelector
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

from settings import DASHSCOPE_API_KEY

# 1-定义客户端
llm = ChatOpenAI(
    model="qwen-max",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    api_key=SecretStr(DASHSCOPE_API_KEY),
    streaming=True
)

# 1. 少样本示例池
examples = [
    {"country": "法国", "capital": "巴黎"},
    {"country": "日本", "capital": "东京"},
    {"country": "巴西", "capital": "巴西利亚"},
]

# 2. 单条示例的 Chat 风格模板
example_prompt = PromptTemplate(
    input_variables=["country", "capital"],
    template="国家：{country} -> 首都：{capital}"
)


# 3.关键：用选择器代替固定 examples
example_selector = LengthBasedExampleSelector(
    examples=examples,
    example_prompt=example_prompt,
    max_length=50,  # 调小一点才能看出“动态丢弃”效果
)

# 3. 构建 FewShotPromptTemplate
few_shot_prompt_template = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="请仿照示例，直接输出答案：",
    suffix="国家：{country} -> 首都：",
    input_variables=["country"],
    example_separator="\n"
)

# 4.1 调用-验证不超过max_length的输入->不会加入问题中（3个示例，1个问题）
print()
chat_prompt_detail = few_shot_prompt_template.format(
    country="加拿大")
print(chat_prompt_detail)

# 4.2添加新示例->不会加入问题中（4个示例，1个问题）
print()
new_example = {"country": "英国", "capital": "伦敦"}
few_shot_prompt_template.example_selector.add_example(new_example)

# 4.3 调用-验证超过max_length的输入->最后一个format才会加入问题中
chat_prompt_detail = few_shot_prompt_template.format(
    country="中国")
print("-----chat_prompt_detail:", chat_prompt_detail)

# 5-打印响应结果
print("-----Hold on, LLM 正在回答！-----")
response = llm.stream(chat_prompt_detail)
for chunk in response:
    # 打印不换行
    print(chunk.content, end="")
